from typing import Tuple
import torch
from loadDataset import MakeDataSetV2
import numpy as np
from torch.utils.data import Dataset


class BCICompDatasetForMulti(Dataset):
    def __init__(self, channelNum: Tuple[int, int, int], subjectARootPath: str, subjectBRootPath: str,
                 needIndex: bool = False, convertToF: bool = True):
        super().__init__()
        self.A_tarDatas, self.A_noTarDatas = MakeDataSetV2.LoadTarAndNoTarData(subjectARootPath, needIndex, convertToF)
        self.B_tarDatas, self.B_noTarDatas = MakeDataSetV2.LoadTarAndNoTarData(subjectBRootPath, needIndex, convertToF)
        A_tarLen = int(np.size(self.A_tarDatas) / np.size(self.A_tarDatas, -1))
        A_noTarLen = int(np.size(self.A_noTarDatas) / np.size(self.A_noTarDatas, -1))
        b_tarLen = int(np.size(self.B_tarDatas) / np.size(self.B_tarDatas, -1))
        b_noTarLen = int(np.size(self.B_noTarDatas) / np.size(self.B_noTarDatas, -1))
        self.dataLen = (A_tarLen + A_noTarLen + b_tarLen + b_noTarLen) / np.size(self.A_tarDatas,
                                                                                 1) / 2  # 这里除以通道数，因为我们固定了通道，除以2是因为有两个人的数据，二我们合并了
        self.datas = np.zeros((int(self.dataLen), 6, 240))
        self.labels = np.zeros(int(self.dataLen), dtype=np.int16)
        count = 0
        for i in range(np.size(self.A_tarDatas, 0)):
            A_electrodes = np.array(
                [self.A_tarDatas[i, channelNum[0], :, :, :], self.A_tarDatas[i, channelNum[1], :, :, :],
                 self.A_tarDatas[i, channelNum[2], :, :, :]])
            B_electrodes = np.array(
                [self.B_tarDatas[i, channelNum[0], :, :, :], self.B_tarDatas[i, channelNum[1], :, :, :],
                 self.B_tarDatas[i, channelNum[2], :, :, :]])
            for j in range(np.size(A_electrodes, 1)):
                for m in range(np.size(A_electrodes, 2)):
                    self.datas[count, 0:3, :] = A_electrodes[:, j, m, :]
                    self.datas[count, 3:6, :] = B_electrodes[:, j, m, :]
                    self.labels[count] = 1
                    count += 1
        for i in range(np.size(self.A_noTarDatas, 0)):
            A_electrodes = np.array(
                [self.A_noTarDatas[i, channelNum[0], :, :, :], self.A_noTarDatas[i, channelNum[1], :, :, :],
                 self.A_noTarDatas[i, channelNum[2], :, :, :]])
            B_electrodes = np.array(
                [self.B_noTarDatas[i, channelNum[0], :, :, :], self.B_noTarDatas[i, channelNum[1], :, :, :],
                 self.B_noTarDatas[i, channelNum[2], :, :, :]])
            for j in range(np.size(A_electrodes, 1)):
                for m in range(np.size(A_electrodes, 2)):
                    self.datas[count, 0:3, :] = A_electrodes[:, j, m, :]
                    self.datas[count, 3:6, :] = B_electrodes[:, j, m, :]
                    self.labels[count] = 0
                    count += 1
        randomArray=np.random.permutation(len(self.datas))
        self.datas = self.datas[randomArray]
        self.labels = self.labels[randomArray]
    def __len__(self) -> int:
        return int(self.dataLen)

    def __getitem__(self, idx) -> Tuple[np.array, np.int16]:
        if idx > self.dataLen:
            # print(idx)
            len = int(self.dataLen)
            return self.datas[len - 1, :, :], self.labels[len - 1]
        return self.datas[idx, :, :], self.labels[idx]


class BCICompTargetDatasetForMulti(Dataset):
    def __init__(self, channelNum: Tuple[int, int, int], subjectARootPath: str, subjectBRootPath: str,
                 needIndex: bool = False, convertToF: bool = True):
        super().__init__()
        self.A_tarDatas, _ = MakeDataSetV2.LoadTarAndNoTarData(subjectARootPath, needIndex, convertToF)
        self.B_tarDatas, _ = MakeDataSetV2.LoadTarAndNoTarData(subjectBRootPath, needIndex, convertToF)
        A_tarLen = int(np.size(self.A_tarDatas) / np.size(self.A_tarDatas, -1))
        b_tarLen = int(np.size(self.B_tarDatas) / np.size(self.B_tarDatas, -1))
        self.dataLen = (A_tarLen + b_tarLen) / np.size(self.A_tarDatas,
                                                       1) / 2  # 这里除以通道数，因为我们固定了通道，除以2是因为有两个人的数据，二我们合并了
        self.datas = np.zeros((int(self.dataLen), 6, 240))
        self.labels = np.zeros(int(self.dataLen), dtype=np.int16)
        count = 0
        for i in range(np.size(self.A_tarDatas, 0)):
            A_electrodes = np.array(
                [self.A_tarDatas[i, channelNum[0], :, :, :], self.A_tarDatas[i, channelNum[1], :, :, :],
                 self.A_tarDatas[i, channelNum[2], :, :, :]])
            B_electrodes = np.array(
                [self.B_tarDatas[i, channelNum[0], :, :, :], self.B_tarDatas[i, channelNum[1], :, :, :],
                 self.B_tarDatas[i, channelNum[2], :, :, :]])
            for j in range(np.size(A_electrodes, 1)):
                for m in range(np.size(A_electrodes, 2)):
                    self.datas[count, 0:3, :] = A_electrodes[:, j, m, :]
                    self.datas[count, 3:6, :] = B_electrodes[:, j, m, :]
                    self.labels[count] = 1
                    count += 1

    def __len__(self) -> int:
        return int(self.dataLen)

    def __getitem__(self, idx) -> Tuple[np.array, np.int16]:
        if idx > self.dataLen:
            # print(idx)
            len = int(self.dataLen)
            return self.datas[len - 1, :, :], self.labels[len - 1]
        return self.datas[idx, :, :], self.labels[idx]


class BCICompNoTargetDatasetForMulti(Dataset):
    def __init__(self, channelNum: Tuple[int, int, int], subjectARootPath: str, subjectBRootPath: str,
                 needIndex: bool = False, convertToF: bool = True):
        super().__init__()
        _, self.A_noTarDatas = MakeDataSetV2.LoadTarAndNoTarData(subjectARootPath, needIndex, convertToF)
        _, self.B_noTarDatas = MakeDataSetV2.LoadTarAndNoTarData(subjectBRootPath, needIndex, convertToF)
        A_noTarLen = int(np.size(self.A_noTarDatas) / np.size(self.A_noTarDatas, -1))
        b_noTarLen = int(np.size(self.B_noTarDatas) / np.size(self.B_noTarDatas, -1))
        self.dataLen = (A_noTarLen + b_noTarLen) / np.size(self.A_tarDatas,
                                                           1) / 2  # 这里除以通道数，因为我们固定了通道，除以2是因为有两个人的数据，二我们合并了
        self.datas = np.zeros((int(self.dataLen), 6, 240))
        self.labels = np.zeros(int(self.dataLen), dtype=np.int16)
        count = 0

        for i in range(np.size(self.A_noTarDatas, 0)):
            A_electrodes = np.array(
                [self.A_noTarDatas[i, channelNum[0], :, :, :], self.A_noTarDatas[i, channelNum[1], :, :, :],
                 self.A_noTarDatas[i, channelNum[2], :, :, :]])
            B_electrodes = np.array(
                [self.B_noTarDatas[i, channelNum[0], :, :, :], self.B_noTarDatas[i, channelNum[1], :, :, :],
                 self.B_noTarDatas[i, channelNum[2], :, :, :]])
            for j in range(np.size(A_electrodes, 1)):
                for m in range(np.size(A_electrodes, 2)):
                    self.datas[count, 0:3, :] = A_electrodes[:, j, m, :]
                    self.datas[count, 3:6, :] = B_electrodes[:, j, m, :]
                    self.labels[count] = 0
                    count += 1

    def __len__(self) -> int:
        return int(self.dataLen)

    def __getitem__(self, idx) -> Tuple[np.array, np.int16]:
        if idx > self.dataLen:
            # print(idx)
            len = int(self.dataLen)
            return self.datas[len - 1, :, :], self.labels[len - 1]
        return self.datas[idx, :, :], self.labels[idx]
